"""
Author: LYh
Description: OK
"""
from speak_out import play_mp3, ask_question
from model import *
from read_cough_data import read_json_file


if __name__ == "__main__":
    model_dir = '../pytorch_model'

    # 选择模型运行设备、架构
    arch, device = get_model()
    if arch == 'pytorch':
        from transformers import T5ForConditionalGeneration, T5Tokenizer
        import torch
        tokenizer = T5Tokenizer.from_pretrained(model_dir, legacy=False, local_files_only=True)

        print('开始加载权重...', end='')
        if device == 'cpu':
            model = T5ForConditionalGeneration.from_pretrained(model_dir, local_files_only=True)
        elif device == 'npu':
            import torch_npu
            from torch_npu.contrib import transfer_to_npu
            model = T5ForConditionalGeneration.from_pretrained(model_dir, device_map='npu:0', local_files_only=True)  # 使用device_map加载的模型无法被import
            device = torch.device(device)
        elif device == 'cuda':
            model = T5ForConditionalGeneration.from_pretrained(model_dir, device_map='cuda:0', local_files_only=True)
            device = torch.device(device)
        print('加载完毕！')
    elif arch == 'paddle':
        raise ValueError('paddle model is not supported yet.')
        import paddle
        # 使用paddlenlp加载tokenizer
        from paddlenlp.transformers import AutoTokenizer, T5ForConditionalGeneration

        device = paddle.device('npu')
        tokenizer = AutoTokenizer.from_pretrained(model_dir, from_hf_hub=False)
        model = T5ForConditionalGeneration.from_pretrained(model_dir, from_hf_hub=False)


    # 正式推理
    history = []
    cough_data = read_json_file('../cough_data/cough.json')
    cough_env = read_json_file('../cough_data/cough_environment.json')
    query = (f'{{症状:{[e["comments"] for e in cough_env["coughing_instances"][-3:]]},'
            f'环境:[{cough_env["external_environment"]["weather"]["comments"], cough_env["external_environment"]["pollution_level"]["comments"]}'
            f'{cough_env["additional_notes"]}]}}。'
            '根据上述记录分析我的生病原因并给出建议，口语化表达，不要打招呼。')
    # print(query)

    history.append(query)
    response = answer(tokenizer=tokenizer, model=model, text=query, device=device)
    history.append(response)

    print(f"bot：嗨，最近发现你经常咳嗽呢。{response}")
    speak_text = f'嗨，最近发现你经常咳嗽呢。{response}'
    play_mp3(text=speak_text)


    # 模拟大模型提问
    dialog = ask_question()
    history.append(dialog)
    response = answer(tokenizer=tokenizer, model=model, text=str(dialog)+'根据之前我回答你的情况对进一步分析我的病情', device=device)
    history.append(response)

    print(f"bot：{response}")
    play_mp3(text=response)

    print('历史记录：', *history, sep='\n')
